#!/usr/bin/env python
#-*- coding:utf-8 -*-

import math
from random import random, randint
from planner.planner import Planner 
from graph.node_tree import NodeTree
from common.node import Node
from common.edge import Edge
from planner.astar_planner import AStarPlanner

class RRTPlanner(Planner):
    """
    暂时参考：https://www.cnblogs.com/21207-iHome/p/7210543.html
    """
    def __init__(self, graph=None, start=None, goal=None):
        super(RRTPlanner, self).__init__(graph, start, goal)
        self.heading_vec_dict = {}
    
    @property
    def random_tree(self):
        if not hasattr(self, "_random_tree"):
            self._random_tree = NodeTree()
        return self._random_tree

    @random_tree.setter
    def random_tree(self, random_tree):
        self._random_tree = random_tree

    def theta(self, vec1, vec2):
        if vec1 == (0, 0):
            return math.atan2(vec2[1], vec2[0])
        if vec2 == (0, 0):
            return math.atan2(vec1[1], vec1[0])
        v1_mod = (vec1[0] * vec1[0] +  vec1[1] * vec1[1]) ** 0.5
        v2_mod = (vec2[0] * vec2[0] +  vec2[1] * vec2[1]) ** 0.5
        v1_v2_in_prod = vec1[0] * vec2[0] + vec1[1] * vec2[1]
        return math.acos(v1_v2_in_prod / (v1_mod * v2_mod + 0.000001))

    def _sample_target(self, random_prob=0.3, block_check=False, max_repeat=10):
        _prob = random()
        if _prob <= random_prob:
            _x = randint(0, self.graph.width - 1)
            _y = randint(0, self.graph.height - 1)
            if not block_check:
                return Node((_x, _y), (_x, _y))
            for _ in range(max_repeat):
                _x = randint(0, self.graph.width - 1)
                _y = randint(0, self.graph.height - 1)
                if not self.blocked((_x, _y)):
                    return Node((_x, _y), (_x, _y))
        return self.goal_node

    def motion_cost(self, from_node, to_node, rate_dist=0.4):
        # print("55", self.heading_vec_dict)
        # print("56", from_node, to_node)
        _target_vec = (to_node.pos[0] - from_node.pos[0], to_node.pos[1] - from_node.pos[1])
        delta_dist = ((to_node.pos[0] - from_node.pos[0]) ** 2 + (to_node.pos[1] - from_node.pos[1]) ** 2 ) ** 0.5
        delta_theta = self.theta(self.heading_vec_dict[from_node.nid], _target_vec)
        cost_of_dist = max(1/(delta_dist + 0.0000001), delta_dist)
        cost_of_theta = abs(math.tan(delta_theta))
        # total_cost = cost_of_dist + cost_of_theta
        total_cost = cost_of_dist * rate_dist + cost_of_theta * (1 - rate_dist)
        return total_cost
            
    def _find_nearest_node(self, target_node, by='dist', leaf=False): #may not valiable
        nearest_node = None
        nearest_val = float('inf')
        for candi_node in self.random_tree.nodes():
            # print("44", type(candi_node), candi_node)
            if leaf and not candi_node.leaf:
                continue
            if by == 'cost':
                val = self.cost(candi_node, target_node)
            if by == 'motion':
                val = self.motion_cost(candi_node, target_node, rate_dist=0.6)
            else: # dist
                # print("50", candi_node, target_node)
                val = self.distance(candi_node, target_node)
            if val < nearest_val:
                nearest_val = val
                nearest_node = candi_node
        return nearest_node, nearest_val

    def _extend_node(self, from_node, to_node, dist= 5, block_check=True):
        from_pos = from_node.pos
        to_pos = to_node.pos
        base_dist = ((to_pos[0] - from_pos[0]) ** 2 + (to_pos[1] - from_pos[1]) ** 2) ** 0.5
        if base_dist < dist / 2:
            return None
        rate = dist / (base_dist + 0.000000001)
        new_x = int(rate * (to_pos[0] - from_pos[0]) + from_pos[0])
        new_y = int(rate * (to_pos[1] - from_pos[1]) + from_pos[1])
        return Node((new_x, new_y), (new_x, new_y))

    def _build_tree(self, random_prob=0.3, step_length=10, max_epoch=200): 
        self.start_node = Node(self.start, self.start)
        self.heading_vec_dict[self.start] = (self.goal[0] - self.start[0], self.goal[1] - self.start[1])
        self.goal_node = Node(self.goal, self.goal)
        self.random_tree.add_node(self.start_node)
        self.random_tree.root = self.start
        for e in range(max_epoch):
            print("epoch: {}".format(e))
            target_node = self._sample_target(random_prob, block_check=False, max_repeat=10) #重复的点
            # nearest_node, nearest_val = self._find_nearest_node(target_node, by='dist')#重复的点
            nearest_node, nearest_val = self._find_nearest_node(target_node, by='motion')#重复的点
            if nearest_val < 1: 
                continue
            new_node = self._extend_node(nearest_node, target_node, step_length)
            if new_node and not self.blocked(nearest_node, new_node):
                self.random_tree.add_node(new_node)
                _weight = self.cost(nearest_node, new_node)
                self.random_tree.add_edge(Edge(nearest_node.nid, new_node.nid, _weight))
                _heading_vec = (new_node.pos[0] - nearest_node.pos[0], new_node.pos[1] - nearest_node.pos[1])
                self.heading_vec_dict[new_node.nid] = _heading_vec
                almost_reach, remaining_dist = self._almost_reach_goal(new_node, threshold=5)
                if almost_reach:
                    if remaining_dist > 0:
                        self.random_tree.add_node(self.goal_node)
                        _weight = self.cost(new_node, self.goal_node)
                        self.random_tree.add_edge(Edge(new_node.nid, self.goal_node.nid, _weight))
                        _heading_vec = (self.goal_node.pos[0] - new_node.pos[0], self.goal_node.pos[1] - new_node.pos[1])
                        self.heading_vec_dict[self.goal] = _heading_vec
                    print("reach goal at epoch: {}".format(e))
                    return True
        print("sampling failed in {} epoch".format(max_epoch))
        return False

    def _almost_reach_goal(self, cur_node, threshold=None):
        if cur_node == self.goal_node:
            return True, 0
        if threshold and not self.blocked(cur_node, self.goal_node) \
            and self.cost(cur_node, self.goal_node) <= threshold:
            return True, self.cost(cur_node, self.goal_node)
        return False, float('inf')

    def _solve_heuristic_to_goal(self):
        heurs = {}
        for node in self.random_tree.nodes():
            heurs[(node.nid, self.goal)] = self.graph.distance(node, self.goal_node)
        return heurs

    def _search_path(self):
        self.random_tree.heuristics = self._solve_heuristic_to_goal()
        _planner = AStarPlanner()
        _planner.plan(self.random_tree, self.start, self.goal)
        self.result = _planner.result

    def plan(self, graph=None, start=None, goal=None, random_prob=0.4, step_length=8, max_epoch=10000):
        print("plan in graph {} from {} to {}".format(graph, start, goal))
        if not self.check(graph, start, goal):
            return None
        if self._build_tree(random_prob=random_prob, step_length=step_length, max_epoch=max_epoch):
            self._search_path()
        return self.result
